#ifndef SWITCH_MMU_H
#define SWITCH_MMU_H

#include <unordered_map>
#include <ns3/node.h>

namespace ns3 {

struct FlowState{
	uint32_t sentBytes;
	uint32_t pkgCnt;			// number of packets in current queue
	uint64_t timestampMs;			// us
	
	FlowState(): sentBytes(0), pkgCnt(0), timestampMs(0) {}
};

// queue level
struct FlowStateTable{
	std::unordered_map<uint32_t, FlowState> table;
	uint32_t activeFlowCnt;
	uint64_t lastUploadTimeMs;
	uint32_t totalSendBytes;

	uint64_t lastPrintTime;
	
	static double alpha;  // static alpha

	FlowStateTable(): lastPrintTime(0), activeFlowCnt(0), lastUploadTimeMs(0), totalSendBytes(0) {}

	void Update(uint32_t flow_hash, uint32_t size){
		uint64_t currentTimeMs = Simulator::Now().GetMilliSeconds();

		if(currentTimeMs - table[flow_hash].timestampMs > 100){		// expire time =  10ms
			table[flow_hash].sentBytes = 0;
			table[flow_hash].pkgCnt = 0;
			// table[flow_hash].timestampMs = currentTimeMs;  // meng teacher
		}

		if(table[flow_hash].pkgCnt == 0){
			activeFlowCnt ++ ;
			totalSendBytes += table[flow_hash].sentBytes;
		}
		totalSendBytes += size;

		table[flow_hash].timestampMs = currentTimeMs;
		table[flow_hash].sentBytes += size;
		table[flow_hash].pkgCnt++;

	}

	void Leave(uint32_t flow_hash){
		table[flow_hash].pkgCnt--;
		if(table[flow_hash].pkgCnt == 0) {
			activeFlowCnt -- ;
			totalSendBytes -= table[flow_hash].sentBytes;
		}
	}

	double GetHPCCAmend(uint32_t flow_hash) {
		if(activeFlowCnt <= 1) return 1;
		if(alpha == 100) return 1; 		// for best alpha ex, using alpha=100 to shut down DiffCC
		uint32_t avgSendBytes = totalSendBytes / activeFlowCnt;
		if(table[flow_hash].sentBytes < avgSendBytes) return alpha;
		else return 2 - alpha;
	}

	double GetDCQCNAmend(uint32_t flow_hash) {
		if(activeFlowCnt <= 1) return 1;
		if(alpha == 100) return 1; 		// for best alpha ex, using alpha=100 to shut down DiffCC
		uint32_t avgSendBytes = totalSendBytes / activeFlowCnt;
		// PrintActive(0);
		// printf("%u\t%u\tdecided:%d\n", flow_hash, avgSendBytes, table[flow_hash].sentBytes < avgSendBytes);
		if(table[flow_hash].sentBytes < avgSendBytes) return 2 - alpha;
		else return alpha;
	}

	void Print(uint32_t flow_hash){
		printf("flow_id: %u\t sentBytes: %.3lfKB\tpkgCnt: %u\n", flow_hash, table[flow_hash].sentBytes *1.0 / 1000, table[flow_hash].pkgCnt);
	}

	int PrintActive(uint64_t gapTimeUs){
		if(Simulator::Now().GetMicroSeconds() - lastPrintTime < gapTimeUs)		// print every 50us
			return 0;
		lastPrintTime = Simulator::Now().GetMicroSeconds();
		int cnt = 0;
		for (auto it = table.begin(); it != table.end(); it++){
			if (it->second.pkgCnt > 0){
				cnt++;
				// Print(it->first);
			}
		}
		
		// for active bytes stat
		if(cnt <= 1) return 0;
		
		for (auto it = table.begin(); it != table.end(); it++){
			if (it->second.pkgCnt > 0){
				cnt++;
				// Print(it->first);
			}
		}
		// printf("Time:%luns\n", Simulator::Now().GetNanoSeconds());
		
		return cnt;
	}
};

class Packet;

class SwitchMmu: public Object{
public:
	static const uint32_t pCnt = 257;	// Number of ports used
	static const uint32_t qCnt = 8;	// Number of queues/priorities used

	static TypeId GetTypeId (void);

	SwitchMmu(void);

	bool CheckIngressAdmission(uint32_t port, uint32_t qIndex, uint32_t psize);
	bool CheckEgressAdmission(uint32_t port, uint32_t qIndex, uint32_t psize);
	void UpdateIngressAdmission(uint32_t port, uint32_t qIndex, uint32_t psize);
	void UpdateEgressAdmission(uint32_t port, uint32_t qIndex, uint32_t psize);
	void RemoveFromIngressAdmission(uint32_t port, uint32_t qIndex, uint32_t psize);
	void RemoveFromEgressAdmission(uint32_t port, uint32_t qIndex, uint32_t psize);

	bool CheckShouldPause(uint32_t port, uint32_t qIndex);
	bool CheckShouldResume(uint32_t port, uint32_t qIndex);
	void SetPause(uint32_t port, uint32_t qIndex);
	void SetResume(uint32_t port, uint32_t qIndex);
	//void GetPauseClasses(uint32_t port, uint32_t qIndex);
	//bool GetResumeClasses(uint32_t port, uint32_t qIndex);

	uint32_t GetPfcThreshold(uint32_t port);
	uint32_t GetSharedUsed(uint32_t port, uint32_t qIndex);

	bool ShouldSendCN(uint32_t ifindex, uint32_t qIndex, uint32_t flow_hash);

	void ConfigEcn(uint32_t port, uint32_t _kmin, uint32_t _kmax, double _pmax);
	void ConfigHdrm(uint32_t port, uint32_t size);
	void ConfigNPort(uint32_t n_port);
	void ConfigBufferSize(uint32_t size);

	// config
	uint32_t node_id;
	uint32_t buffer_size;
	uint32_t pfc_a_shift[pCnt];
	uint32_t reserve;
	uint32_t headroom[pCnt];
	uint32_t resume_offset;
	uint32_t kmin[pCnt], kmax[pCnt];
	double pmax[pCnt];
	uint32_t total_hdrm;
	uint32_t total_rsrv;

	// runtime
	uint32_t shared_used_bytes;
	uint32_t hdrm_bytes[pCnt][qCnt];
	uint32_t ingress_bytes[pCnt][qCnt];
	uint32_t paused[pCnt][qCnt];
	uint32_t egress_bytes[pCnt][qCnt];

	FlowStateTable flow_table[pCnt];
};

} /* namespace ns3 */

#endif /* SWITCH_MMU_H */

